#!/usr/bin/env python3

# Copyright (C) 2025- Kui, BI8DHC
#
# Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer
#     in the documentation and/or other materials provided with the distribution.
# 
#  3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived
#     from this software without specific prior written permission.
# 
#  4. Use of this source code or binary releases for commercial purposes is strictly forbidden. This includes, without limitation,
#     incorporation in a commercial product or incorporation into a product or project which allows commercial use.
# 
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
#  HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
#  ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
# USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


import sys
import os
import argparse
import time
import serial
import logging
from tqdm import tqdm

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("flash_operations.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger()

# Flash constants
FLASH_PAGE_SIZE = 256
FLASH_SECTOR_SIZE = 4096
W25Q128_SIZE = 16 * 1024 * 1024  # 16MB

# Command definitions from usb_com.c
CPS_ACCESS_FLASH = 1
CPS_ACCESS_EEPROM = 2
CPS_ACCESS_MCU_ROM = 5
CPS_ACCESS_DISPLAY_BUFFER = 6

class FlashTool:
    def __init__(self, port, baudrate=115200, timeout=1):
        self.port = port
        self.baudrate = baudrate
        self.timeout = timeout
        self.ser = None
        self.connected = False

    def connect(self):
        """Connect to the COM port"""
        try:
            self.ser = serial.Serial(
                port=self.port,
                baudrate=self.baudrate,
                timeout=self.timeout
            )
            self.connected = True
            logger.info(f"Connected to {self.port} at {self.baudrate} baud")
            return True
        except serial.SerialException as e:
            logger.error(f"Failed to connect to {self.port}: {e}")
            return False

    def disconnect(self):
        """Disconnect from the COM port"""
        if self.ser and self.ser.is_open:
            self.ser.close()
            self.connected = False
            logger.info(f"Disconnected from {self.port}")

    def send_command(self, cmd_data, expected_response_length=None):
        """Send a command and read the response"""
        if not self.connected:
            logger.error("Not connected to device")
            return None

        try:
            self.ser.write(cmd_data)
            
            # Read the first byte to determine if operation was successful
            response_header = self.ser.read(1)
            if not response_header:
                logger.error("No response received")
                return None

            if response_header == b'-':
                logger.error("Command failed")
                return None

            if expected_response_length is None:
                # For commands that don't have a fixed response length
                response = self.ser.read(self.ser.in_waiting)
                return response_header + response
            else:
                # For read operations where we expect a known amount of data
                response_len_bytes = self.ser.read(2)
                if len(response_len_bytes) != 2:
                    logger.error("Failed to read response length")
                    return None
                
                response_length = (response_len_bytes[0] << 8) | response_len_bytes[1]
                response_data = self.ser.read(response_length)
                
                if len(response_data) != response_length:
                    logger.error(f"Incomplete response. Expected {response_length} bytes, got {len(response_data)}")
                    return None
                
                return response_header + response_len_bytes + response_data
        except serial.SerialException as e:
            logger.error(f"Serial communication error: {e}")
            return None

    def prepare_sector(self, sector_number):
        """Prepare a sector for writing"""
        cmd = bytearray([ord('X'), 1])  # Flash Prepare Sector command
        cmd.extend([(sector_number >> 16) & 0xFF,
                   (sector_number >> 8) & 0xFF,
                   sector_number & 0xFF])
        
        response = self.send_command(cmd)
        return response and response[0] == ord('X')

    def send_data(self, address, data):
        """Send data to be written"""
        length = len(data)
        cmd = bytearray([ord('X'), 2])  # Flash Send Data command
        cmd.extend([(address >> 24) & 0xFF,
                   (address >> 16) & 0xFF,
                   (address >> 8) & 0xFF,
                   address & 0xFF])
        cmd.extend([(length >> 8) & 0xFF, length & 0xFF])
        cmd.extend(data)
        
        response = self.send_command(cmd)
        return response and response[0] == ord('X')

    def write_sector(self):
        """Commit the write operation"""
        cmd = bytearray([ord('X'), 3])  # Flash Write command
        
        response = self.send_command(cmd)
        return response and response[0] == ord('X')

    def write_flash(self, start_address, data):
        """Write data to flash memory"""
        if not self.connected:
            if not self.connect():
                return False

        # Make sure start address is aligned to sector boundary
        sector_start = (start_address // FLASH_SECTOR_SIZE) * FLASH_SECTOR_SIZE
        offset_in_sector = start_address - sector_start
        
        # Calculate how many sectors we need to write
        total_length = len(data)
        end_address = start_address + total_length
        end_sector = (end_address + FLASH_SECTOR_SIZE - 1) // FLASH_SECTOR_SIZE
        start_sector = sector_start // FLASH_SECTOR_SIZE
        
        logger.info(f"Writing {total_length} bytes starting at address 0x{start_address:08X}")
        logger.info(f"This will affect sectors {start_sector} to {end_sector-1}")
        
        bytes_written = 0
        current_addr = start_address
        
        # Process each sector
        with tqdm(total=total_length, unit='B', unit_scale=True, desc="Writing") as pbar:
            for sector in range(start_sector, end_sector):
                sector_addr = sector * FLASH_SECTOR_SIZE
                
                # Read the current sector content if needed (for first or last sector)
                if sector == start_sector and offset_in_sector > 0:
                    logger.info(f"Reading sector {sector} for partial update")
                    sector_data = self.read_flash(sector_addr, FLASH_SECTOR_SIZE)
                    if sector_data is None:
                        logger.error(f"Failed to read sector {sector}")
                        return False
                    
                    # Prepare the new sector data
                    bytes_to_write = min(FLASH_SECTOR_SIZE - offset_in_sector, total_length - bytes_written)
                    new_sector_data = bytearray(sector_data)
                    new_sector_data[offset_in_sector:offset_in_sector+bytes_to_write] = data[bytes_written:bytes_written+bytes_to_write]
                                       
                else:
                    # Full sector write
                    bytes_to_write = min(FLASH_SECTOR_SIZE, total_length - bytes_written)
                    new_sector_data = data[bytes_written:bytes_written+bytes_to_write]
                    # If less than a full sector, pad with erased state (0xFF)
                    if bytes_to_write < FLASH_SECTOR_SIZE:
                        new_sector_data = new_sector_data + bytearray([0xFF] * (FLASH_SECTOR_SIZE - bytes_to_write))
                
                # Prepare the sector
                logger.info(f"Preparing sector {sector}")
                if not self.prepare_sector(sector):
                    logger.error(f"Failed to prepare sector {sector}")
                    return False
                
                # Send data in chunks (pages)
                for page in range(0, FLASH_SECTOR_SIZE, FLASH_PAGE_SIZE):
                    page_data = new_sector_data[page:page+FLASH_PAGE_SIZE]
                    page_addr = sector_addr + page
                    
                    logger.debug(f"Sending data for page at 0x{page_addr:08X}")
                    if not self.send_data(page_addr, page_data):
                        logger.error(f"Failed to send data for page at 0x{page_addr:08X}")
                        return False
                
                # Write the sector
                logger.info(f"Writing sector {sector}")
                if not self.write_sector():
                    logger.error(f"Failed to write sector {sector}")
                    return False
                
                # Update progress
                if sector == start_sector:
                    if offset_in_sector > 0:
                        actual_bytes_written = min(FLASH_SECTOR_SIZE - offset_in_sector, total_length)
                    else:
                        actual_bytes_written = min(FLASH_SECTOR_SIZE, total_length)
                elif sector == end_sector - 1:
                    actual_bytes_written = (end_address % FLASH_SECTOR_SIZE) or FLASH_SECTOR_SIZE
                else:
                    actual_bytes_written = FLASH_SECTOR_SIZE
                
                bytes_written += actual_bytes_written
                pbar.update(actual_bytes_written)
        
        logger.info(f"Successfully wrote {bytes_written} bytes")
        return True

    def write_file_to_flash(self, start_address, input_file):
        """Write file content to flash memory"""
        try:
            file_size = os.path.getsize(input_file)
            logger.info(f"Reading {file_size} bytes from {input_file}")
            
            with open(input_file, 'rb') as f:
                data = f.read()
            
            return self.write_flash(start_address, data)
            
        except IOError as e:
            logger.error(f"File operation error: {e}")
            return False

def parse_address(addr_str):
    """Parse address string in decimal or hex format"""
    try:
        if addr_str.startswith("0x"):
            return int(addr_str, 16)
        else:
            return int(addr_str)
    except ValueError:
        logger.error(f"Invalid address format: {addr_str}")
        sys.exit(1)


def main():
    parser = argparse.ArgumentParser(description='Font Flash Writer Tool')
    parser.add_argument('--port', required=True, help='COM port to use (e.g., COM3 or /dev/ttyUSB0)')
    
    subparsers = parser.add_subparsers(dest='command', help='Command')
      
    args = parser.parse_args()

    input = "font.bin"              # large font
    input_s = "font_s.bin"            # small font
    #start_address = "0xFE0000"
    large_start_address = "0xFD0000"      # large font start address
    small_start_address = "0xFF0000"      # small font start address
    
    # Validate the COM port
    if not args.port:
        logger.error("COM port is required")
        parser.print_help()
        sys.exit(1)
    
    # Create flash tool instance
    flash_tool = FlashTool(args.port)
    
    try:
        if not flash_tool.connect():
            sys.exit(1)
        
        # large font upload
        start_address = parse_address(large_start_address)                      
        
        logger.info(f"Writing {input} to flash at 0x{start_address:X}")
        success = flash_tool.write_file_to_flash(start_address, input)
        
        if success:
            logger.info("Upload large font completed successfully")
            
        else:
            logger.error("Upload large font failed")
            sys.exit(1)
            
        # small font upload
        start_address = parse_address(small_start_address)                      
        
        logger.info(f"Writing  {input_s} to flash at 0x{start_address:X}")
        success = flash_tool.write_file_to_flash(start_address, input_s)
        
        if success:
            logger.info("Upload small font successfully")
            
        else:
            logger.error("Upload small font failed")
            sys.exit(1)
                
    finally:
        flash_tool.disconnect()


if __name__ == "__main__":
    main()